package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/hashicorp/golang-lru/v2/expirable"
	"golang.org/x/net/netutil"
	"golang.org/x/time/rate"
)

// 项目版本（可通过 -ldflags="-X main.buildVersion=..." 覆盖）
var buildVersion = "EchoIP/v0.0.1"

var (
	endpoints = []string{"/ip", "/healthz", "/version", "/metrics"}
	reqTotal  = make(map[string]*uint64)
	errTotal  = make(map[string]*uint64)
)

var trustedNets []*net.IPNet
var ipLimiters *expirable.LRU[string, *rate.Limiter]

const (
	ipLimiterBurst     = 5
	ipLimiterRate      = 10 * time.Second
	globalLimiterBurst = 20
	globalLimiterRate  = 2 * time.Second
)

func init() {
	for _, ep := range endpoints {
		var r uint64
		reqTotal[ep] = &r
		var e uint64
		errTotal[ep] = &e
	}

	for _, cidrStr := range []string{
		"127.0.0.0/8",
		"::1/128",
		"10.0.0.0/8",
		"172.16.0.0/12",
		"192.168.0.0/16",
	} {
		_, netw, err := net.ParseCIDR(cidrStr)
		if err == nil {
			trustedNets = append(trustedNets, netw)
		}
	}

	ipLimiters = expirable.NewLRU[string, *rate.Limiter](1000, nil, 5*time.Minute)
	log.SetFlags(log.LstdFlags | log.Lmicroseconds | log.LUTC)
}

func incReq(endpoint string) { atomic.AddUint64(reqTotal[endpoint], 1) }
func incErr(endpoint string) { atomic.AddUint64(errTotal[endpoint], 1) }

func isTrustedIP(ip net.IP) bool {
	if ip == nil {
		return false
	}
	for _, netw := range trustedNets {
		if netw.Contains(ip) {
			return true
		}
	}
	return false
}

func isValidIP(s string) bool {
	s = strings.TrimSpace(s)
	if s == "" {
		return false
	}
	ip := net.ParseIP(s)
	return ip != nil && (ip.To4() != nil || ip.To16() != nil)
}

func getClientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		host = r.RemoteAddr
	}
	clientIP := net.ParseIP(host)
	if clientIP == nil {
		return host
	}

	if isTrustedIP(clientIP) {
		if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
			for _, ipStr := range strings.Split(xff, ",") {
				if isValidIP(ipStr) {
					return strings.TrimSpace(ipStr)
				}
			}
		}
		if xri := r.Header.Get("X-Real-IP"); isValidIP(xri) {
			return strings.TrimSpace(xri)
		}
	}

	return clientIP.String()
}

func setSecurityHeaders(w http.ResponseWriter) {
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("X-Frame-Options", "DENY")
	w.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate, private")
	w.Header().Set("Pragma", "no-cache")
}

func isValidHost(host string) bool {
	if host == "" {
		return true
	}
	host = strings.Split(strings.ToLower(host), ":")[0]
	return host == "localhost" || host == "127.0.0.1" || host == "::1"
}

func getIPLimiter(ip string) *rate.Limiter {
	limiter, exists := ipLimiters.Get(ip)
	if !exists {
		limiter = rate.NewLimiter(rate.Every(ipLimiterRate), ipLimiterBurst)
		ipLimiters.Add(ip, limiter)
	}
	return limiter
}

var globalLimiter = rate.NewLimiter(rate.Every(globalLimiterRate), globalLimiterBurst)

func makeHandler(endpoint string, disableLogging bool, handlerFunc func(http.ResponseWriter, *http.Request)) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		incReq(endpoint)

		if !globalLimiter.Allow() {
			incErr(endpoint)
			http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
			return
		}

		clientIP := getClientIP(r)
		if !getIPLimiter(clientIP).Allow() {
			incErr(endpoint)
			http.Error(w, "Too Many Requests from your IP", http.StatusTooManyRequests)
			return
		}

		if r.Method != http.MethodGet {
			incErr(endpoint)
			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}

		if !isValidHost(r.Host) {
			incErr(endpoint)
			http.Error(w, "Invalid Host header", http.StatusBadRequest)
			return
		}

		handlerFunc(w, r)
	}
}

func ipHandler(disableLogging bool) func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		ip := getClientIP(r)
		if !disableLogging {
			if cn := r.Header.Get("X-SSL-Client-CN"); cn != "" {
				log.Printf("[EchoIP] Client [%s] requested IP: %s", cn, ip)
			} else {
				log.Printf("[EchoIP] Client requested IP: %s", ip)
			}
		}
		setSecurityHeaders(w)
		fmt.Fprintln(w, ip)
	}
}

func healthHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		setSecurityHeaders(w)
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("OK\n"))
	}
}

func versionHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		setSecurityHeaders(w)
		fmt.Fprintf(w, "%s\n", buildVersion)
	}
}

func metricsHandler() func(http.ResponseWriter, *http.Request) {
	return func(w http.ResponseWriter, r *http.Request) {
		setSecurityHeaders(w)
		w.Header().Set("Content-Type", "text/plain; version=0.0.4")
		for _, ep := range endpoints {
			fmt.Fprintf(w, "# HELP echoip_requests_total Total requests to %s\n", ep)
			fmt.Fprintf(w, "# TYPE echoip_requests_total counter\n")
			fmt.Fprintf(w, "echoip_requests_total{endpoint=\"%s\"} %d\n", ep, atomic.LoadUint64(reqTotal[ep]))
			fmt.Fprintf(w, "# HELP echoip_errors_total Errors on %s\n", ep)
			fmt.Fprintf(w, "# TYPE echoip_errors_total counter\n")
			fmt.Fprintf(w, "echoip_errors_total{endpoint=\"%s\"} %d\n", ep, atomic.LoadUint64(errTotal[ep]))
		}
	}
}

func validatePort(portStr string) (int, error) {
	if portStr == "" {
		return 0, fmt.Errorf("port is required and cannot be empty")
	}
	port, err := strconv.Atoi(portStr)
	if err != nil {
		return 0, fmt.Errorf("port must be a number")
	}
	if port < 1024 || port > 65535 {
		return 0, fmt.Errorf("port must be between 1024 and 65535")
	}
	return port, nil
}

func main() {
	var (
		portStr        string
		disableLogging bool
	)
	flag.StringVar(&portStr, "port", "", "REQUIRED: local listening port (1024-65535)")
	flag.BoolVar(&disableLogging, "disable-logging", false, "Disable client request logging")
	flag.Parse()

	if portStr == "" {
		fmt.Fprintln(os.Stderr, "Error: -port flag is required")
		flag.Usage()
		os.Exit(1)
	}

	port, err := validatePort(portStr)
	if err != nil {
		log.Fatalf("[EchoIP] Invalid port: %v", err)
	}

	addr := fmt.Sprintf("127.0.0.1:%d", port)

	http.HandleFunc("/ip", makeHandler("/ip", disableLogging, ipHandler(disableLogging)))
	http.HandleFunc("/healthz", makeHandler("/healthz", false, healthHandler()))
	http.HandleFunc("/version", makeHandler("/version", false, versionHandler()))
	http.HandleFunc("/metrics", makeHandler("/metrics", false, metricsHandler()))

	listener, err := net.Listen("tcp", addr)
	if err != nil {
		log.Fatal(err)
	}
	defer listener.Close()
	limitListener := netutil.LimitListener(listener, 100)

	server := &http.Server{
		Addr:              addr,
		ReadTimeout:       5 * time.Second,
		WriteTimeout:      5 * time.Second,
		IdleTimeout:       60 * time.Second,
		ReadHeaderTimeout: 5 * time.Second,
	}

	go func() {
		log.Printf("[EchoIP] Starting server on %s (version: %s)", addr, buildVersion)
		if err := server.Serve(limitListener); err != nil && err != http.ErrServerClosed {
			log.Fatalf("[EchoIP] Server failed: %v", err)
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	log.Println("[EchoIP] Shutdown signal received, stopping...")

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := server.Shutdown(ctx); err != nil {
		log.Fatalf("[EchoIP] Forced shutdown: %v", err)
	}

	listener.Close()
	log.Println("[EchoIP] Exited gracefully")
}
